{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4a033743",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([-0.35493255,  0.93489194, -0.32008162], dtype=float32)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import gym\n",
    "\n",
    "\n",
    "#定义环境\n",
    "class MyWrapper(gym.Wrapper):\n",
    "    def __init__(self):\n",
    "        env = gym.make('Pendulum-v1', render_mode='rgb_array')\n",
    "        super().__init__(env)\n",
    "        self.env = env\n",
    "        self.step_n = 0\n",
    "\n",
    "    def reset(self):\n",
    "        state, _ = self.env.reset()\n",
    "        self.step_n = 0\n",
    "        return state\n",
    "\n",
    "    def step(self, action):\n",
    "        state, reward, terminated, truncated, info = self.env.step(action)\n",
    "        done = terminated or truncated\n",
    "        self.step_n += 1\n",
    "        if self.step_n >= 200:\n",
    "            done = True\n",
    "        return state, reward, done, info\n",
    "\n",
    "\n",
    "env = MyWrapper()\n",
    "\n",
    "env.reset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9c383bd4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(200,\n",
       " ([-0.9102440476417542, -0.4140722155570984, 0.9576234817504883],\n",
       "  0.35440483689308167,\n",
       "  -7.461259538753717,\n",
       "  [-0.8951918482780457, -0.44568097591400146, 0.700230062007904],\n",
       "  False),\n",
       " torch.Size([200, 4]),\n",
       " torch.Size([200, 4]))"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "\n",
    "class Pool:\n",
    "    def __init__(self, limit):\n",
    "        #样本池\n",
    "        self.datas = []\n",
    "        self.limit = limit\n",
    "\n",
    "    def add(self, state, action, reward, next_state, over):\n",
    "        if isinstance(state, np.ndarray) or isinstance(state, torch.Tensor):\n",
    "            state = state.reshape(3).tolist()\n",
    "\n",
    "        action = float(action)\n",
    "\n",
    "        reward = float(reward)\n",
    "\n",
    "        if isinstance(next_state, np.ndarray) or isinstance(\n",
    "                next_state, torch.Tensor):\n",
    "            next_state = next_state.reshape(3).tolist()\n",
    "\n",
    "        over = bool(over)\n",
    "\n",
    "        self.datas.append((state, action, reward, next_state, over))\n",
    "        #数据上限,超出时从最古老的开始删除\n",
    "        while len(self.datas) > self.limit:\n",
    "            self.datas.pop(0)\n",
    "\n",
    "    #获取一批数据样本\n",
    "    def get_sample(self):\n",
    "        #从样本池中采样\n",
    "        samples = self.datas\n",
    "\n",
    "        #[b, 3]\n",
    "        state = torch.FloatTensor([i[0] for i in samples]).reshape(-1, 3)\n",
    "        #[b, 1]\n",
    "        action = torch.FloatTensor([i[1] for i in samples]).reshape(-1, 1)\n",
    "        #[b, 1]\n",
    "        reward = torch.FloatTensor([i[2] for i in samples]).reshape(-1, 1)\n",
    "        #[b, 3]\n",
    "        next_state = torch.FloatTensor([i[3] for i in samples]).reshape(-1, 3)\n",
    "        #[b, 1]\n",
    "        over = torch.LongTensor([i[4] for i in samples]).reshape(-1, 1)\n",
    "\n",
    "        #[b, 4]\n",
    "        input = torch.cat([state, action], dim=1)\n",
    "        #[b, 4]\n",
    "        label = torch.cat([reward, next_state - state], dim=1)\n",
    "\n",
    "        return input, label\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.datas)\n",
    "\n",
    "\n",
    "pool = Pool(100000)\n",
    "\n",
    "\n",
    "#初始化一局游戏的数据\n",
    "def _():\n",
    "    #初始化游戏\n",
    "    state = env.reset()\n",
    "\n",
    "    #玩到游戏结束为止\n",
    "    over = False\n",
    "    while not over:\n",
    "        #随机一个动作\n",
    "        action = env.action_space.sample()[0]\n",
    "\n",
    "        #执行动作,得到反馈\n",
    "        next_state, reward, over, _ = env.step([action])\n",
    "\n",
    "        #记录数据样本\n",
    "        pool.add(state, action, reward, next_state, over)\n",
    "\n",
    "        #更新游戏状态,开始下一个动作\n",
    "        state = next_state\n",
    "\n",
    "\n",
    "_()\n",
    "\n",
    "a, b = pool.get_sample()\n",
    "\n",
    "len(pool), pool.datas[0], a.shape, b.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e64c315e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([5, 64, 4]), torch.Size([5, 64, 4]))"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import random\n",
    "\n",
    "\n",
    "#定义主模型\n",
    "class Model(torch.nn.Module):\n",
    "\n",
    "    #swish激活函数\n",
    "    class Swish(torch.nn.Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "\n",
    "        def forward(self, x):\n",
    "            return x * torch.sigmoid(x)\n",
    "\n",
    "    #定义一个工具层\n",
    "    class FCLayer(torch.nn.Module):\n",
    "        def __init__(self, in_size, out_size):\n",
    "            super().__init__()\n",
    "            self.in_size = in_size\n",
    "\n",
    "            #初始化参数\n",
    "            std = in_size**0.5\n",
    "            std *= 2\n",
    "            std = 1 / std\n",
    "\n",
    "            weight = torch.empty(5, in_size, out_size)\n",
    "            torch.nn.init.normal_(weight, mean=0.0, std=std)\n",
    "\n",
    "            #[5, in, out]\n",
    "            self.weight = torch.nn.Parameter(weight)\n",
    "\n",
    "            #[5, 1, out]\n",
    "            self.bias = torch.nn.Parameter(torch.zeros(5, 1, out_size))\n",
    "\n",
    "        def forward(self, x):\n",
    "            #x -> [5, b, in]\n",
    "\n",
    "            #[5, b, in] * [5, in, out] -> [5, b, out]\n",
    "            x = torch.bmm(x, self.weight)\n",
    "\n",
    "            #[5, b, out] + [5, 1, out] -> [5, b, out]\n",
    "            x = x + self.bias\n",
    "\n",
    "            return x\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        self.sequential = torch.nn.Sequential(\n",
    "            self.FCLayer(4, 200),\n",
    "            self.Swish(),\n",
    "            self.FCLayer(200, 200),\n",
    "            self.Swish(),\n",
    "            self.FCLayer(200, 200),\n",
    "            self.Swish(),\n",
    "            self.FCLayer(200, 200),\n",
    "            self.Swish(),\n",
    "            self.FCLayer(200, 8),\n",
    "            torch.nn.Identity(),\n",
    "        )\n",
    "\n",
    "        self.softplus = torch.nn.Softplus()\n",
    "\n",
    "        self.optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)\n",
    "\n",
    "    def forward(self, x):\n",
    "        #x -> [5, b, 4]\n",
    "\n",
    "        #[5, b, 4] -> [5, b, 8]\n",
    "        x = self.sequential(x)\n",
    "\n",
    "        #[5, b, 8] -> [5, b, 4]\n",
    "        mean = x[..., :4]\n",
    "\n",
    "        #[5, b, 8] -> [5, b, 4]\n",
    "        logvar = x[..., 4:]\n",
    "\n",
    "        #[1, 1, 4] - [5, b, 4] -> [5, b, 4]\n",
    "        logvar = 0.5 - logvar\n",
    "\n",
    "        #[1, 1, 4] - [5, b, 4] -> [5, b, 4]\n",
    "        logvar = 0.5 - self.softplus(logvar)\n",
    "\n",
    "        #[5, b, 4] - [1, 1, 4] -> [5, b, 4]\n",
    "        logvar = logvar + 10\n",
    "\n",
    "        #[5, b, 4] + [1, 1, 4] -> [5, b, 4]\n",
    "        logvar = self.softplus(logvar) - 10\n",
    "\n",
    "        #[5, b, 4],[5, b, 4]\n",
    "        return mean, logvar\n",
    "\n",
    "    def train(self, input, label):\n",
    "        #input -> [b, 4]\n",
    "        #label -> [b, 4]\n",
    "\n",
    "        #反复训练N次\n",
    "        for _ in range(len(input) // 64 * 20):\n",
    "            #从全量数据中抽样64个,反复抽5遍,形成5份数据\n",
    "            #[5, 64]\n",
    "            select = [torch.randperm(len(input))[:64] for _ in range(5)]\n",
    "            select = torch.stack(select)\n",
    "            #[5, b, 4],[5, b, 4]\n",
    "            input_select = input[select]\n",
    "            label_select = label[select]\n",
    "            del select\n",
    "\n",
    "            #模型计算\n",
    "            #[5, b, 4] -> [5, b, 4],[5, b, 4]\n",
    "            mean, logvar = model(input_select)\n",
    "\n",
    "            #计算loss\n",
    "            #[b, 4] - [b, 4] * [b, 4] -> [b, 4]\n",
    "            mse_loss = (mean - label_select)**2 * (-logvar).exp()\n",
    "\n",
    "            #[b, 4] -> [b] -> scala\n",
    "            mse_loss = mse_loss.mean(dim=1).mean()\n",
    "\n",
    "            #[b, 4] -> [b] -> scala\n",
    "            var_loss = logvar.mean(dim=1).mean()\n",
    "\n",
    "            loss = mse_loss + var_loss\n",
    "\n",
    "            self.optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            self.optimizer.step()\n",
    "\n",
    "\n",
    "model = Model()\n",
    "#model.train(torch.randn(200, 4), torch.randn(200, 4))\n",
    "\n",
    "a, b = model(torch.randn(5, 64, 4))\n",
    "a.shape, b.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1bced33b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([200, 1]) torch.Size([200, 3])\n",
      "torch.Size([25])\n"
     ]
    }
   ],
   "source": [
    "class MPC:\n",
    "    def _fake_step(self, state, action):\n",
    "        #state -> [b, 3]\n",
    "        #action -> [b, 1]\n",
    "\n",
    "        #[b, 4]\n",
    "        input = torch.cat([state, action], dim=1)\n",
    "\n",
    "        #重复5遍\n",
    "        #[b, 4] -> [1, b, 4] -> [5, b, 4]\n",
    "        input = input.unsqueeze(dim=0).repeat([5, 1, 1])\n",
    "\n",
    "        #模型计算\n",
    "        #[5, b, 4] -> [5, b, 4],[5, b, 4]\n",
    "        with torch.no_grad():\n",
    "            mean, std = model(input)\n",
    "        std = std.exp().sqrt()\n",
    "        del input\n",
    "\n",
    "        #means的后3列加上环境数据\n",
    "        mean[:, :, 1:] += state\n",
    "\n",
    "        #重采样\n",
    "        #[5, b ,4]\n",
    "        sample = torch.distributions.Normal(0, 1).sample(mean.shape)\n",
    "        sample = mean + sample * std\n",
    "\n",
    "        #0-4的值域采样b个元素\n",
    "        #[4, 4, 2, 4, 3, 4, 1, 3, 3, 0, 2,...\n",
    "        select = [random.choice(range(5)) for _ in range(mean.shape[1])]\n",
    "\n",
    "        #重采样结果,的结果,第0个维度,0-4随机选择,第二个维度,0-b顺序选择\n",
    "        #[5, b ,4] -> [b, 4]\n",
    "        sample = sample[select, range(mean.shape[1])]\n",
    "\n",
    "        #切分一下,就成了rewards, next_state\n",
    "        reward, next_state = sample[:, :1], sample[:, 1:]\n",
    "\n",
    "        return reward, next_state\n",
    "\n",
    "    def _cem_optimize(self, state, mean):\n",
    "        state = torch.FloatTensor(state).reshape(1, 3)\n",
    "        var = torch.ones(25)\n",
    "        #state -> [1, 3]\n",
    "        #mean -> [25]\n",
    "\n",
    "        #当前游戏的环境信息,复制50次\n",
    "        #[1, 3] -> [50, 3]\n",
    "        state = state.repeat(50, 1)\n",
    "\n",
    "        #循环5次,寻找最优解\n",
    "        for _ in range(5):\n",
    "            #采样50个标准正态分布数据作为action\n",
    "            actions = torch.distributions.Normal(0, 1).sample([50, 25])\n",
    "\n",
    "            #乘以标准差,加上均值\n",
    "            #[50, 25] * [25] -> [50, 25]\n",
    "            actions *= var**0.5\n",
    "            #[50, 25] * [25] -> [50, 25]\n",
    "            actions += mean\n",
    "\n",
    "            #计算每条动作序列的累积奖励\n",
    "            #[50, 1]\n",
    "            reward_sum = torch.zeros(50, 1)\n",
    "\n",
    "            #遍历25个动作\n",
    "            for i in range(25):\n",
    "                #[50, 25] -> [50, 1]\n",
    "                action = actions[:, i].unsqueeze(dim=1)\n",
    "\n",
    "                #现在是不能真的去玩游戏的,只能去预测reward和next_state\n",
    "                #[50, 3],[50, 1] -> [50, 1],[50, 3]\n",
    "                reward, state = self._fake_step(state, action)\n",
    "\n",
    "                #[50, 1] + [50, 1] -> [50, 1]\n",
    "                reward_sum += reward\n",
    "\n",
    "            #按照reward_sum从小到大排列\n",
    "            #[50]\n",
    "            select = torch.sort(reward_sum.squeeze(dim=1)).indices\n",
    "            #[50, 25]\n",
    "            actions = actions[select]\n",
    "            del select\n",
    "\n",
    "            #取反馈最优的10个动作链\n",
    "            #[10, 25]\n",
    "            actions = actions[-10:]\n",
    "\n",
    "            #在下一次随机动作时,希望贴近这些动作的分布\n",
    "            #[25]\n",
    "            new_mean = actions.mean(dim=0)\n",
    "            new_var = actions.var(dim=0)\n",
    "\n",
    "            #增量更新\n",
    "            #[25] + [25] -> [25]\n",
    "            mean = 0.1 * mean + 0.9 * new_mean\n",
    "            #[25] + [25] -> [25]\n",
    "            var = 0.1 * var + 0.9 * new_var\n",
    "\n",
    "        return mean\n",
    "\n",
    "    def mpc(self):\n",
    "        #初始化动作的分布均值都是0\n",
    "        mean = torch.zeros(25)\n",
    "\n",
    "        reward_sum = 0\n",
    "        state = env.reset()\n",
    "        over = False\n",
    "        while not over:\n",
    "            #在当前状态下,找25个最优动作的均值\n",
    "            #[1, 3],[25],[25] -> [25]\n",
    "            actions = self._cem_optimize(state, mean)\n",
    "\n",
    "            #执行第一个动作\n",
    "            action = actions[0].item()\n",
    "\n",
    "            #执行动作\n",
    "            next_state, reward, over, _ = env.step([action])\n",
    "\n",
    "            #增加数据\n",
    "            pool.add(state, action, reward, next_state, over)\n",
    "\n",
    "            state = next_state\n",
    "            reward_sum += reward\n",
    "\n",
    "            #下个动作的均值,在当前动作均值的基础上寻找\n",
    "            #[25]\n",
    "            mean = torch.empty(actions.shape)\n",
    "            mean[:-1] = actions[1:]\n",
    "            mean[-1] = 0\n",
    "\n",
    "        return reward_sum\n",
    "\n",
    "\n",
    "mpc = MPC()\n",
    "\n",
    "a, b = mpc._fake_step(torch.randn(200, 3), torch.randn(200, 1))\n",
    "\n",
    "print(a.shape, b.shape)\n",
    "\n",
    "print(mpc._cem_optimize(torch.randn(1, 3), torch.zeros(25)).shape)\n",
    "\n",
    "#print(mpc.mpc())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b748ed1d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 400 -970.1462027358459\n",
      "1 600 -1044.89966441498\n",
      "2 800 -984.9260862715894\n",
      "3 1000 -1691.3082241692753\n",
      "4 1200 -886.7489707469861\n",
      "5 1400 -872.9982830534415\n",
      "6 1600 -1320.1991576977232\n",
      "7 1800 -772.4313091943403\n",
      "8 2000 -844.8528234513566\n",
      "9 2200 -502.5910743767378\n"
     ]
    }
   ],
   "source": [
    "for i in range(10):\n",
    "    input, label = pool.get_sample()\n",
    "    model.train(input, label)\n",
    "    reward_sum = mpc.mpc()\n",
    "    print(i, len(pool), reward_sum)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:pt39]",
   "language": "python",
   "name": "conda-env-pt39-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
